import typing
import albumentations as A
from pathlib import Path
import cv2
from superstream import Stream
import numpy as np
from concurrent.futures import ThreadPoolExecutor as Exec, Future
from functools import cache


from .common import image_suffix


def define_transformer(**kwargs):
    add_blur = A.OneOf((A.MotionBlur(), A.MedianBlur(blur_limit=3), A.Blur(blur_limit=3)))
    random_shift = A.OneOf(
        (
            A.RandomBrightnessContrast(),
            A.ShiftScaleRotate(
                shift_limit=0.625,
                scale_limit=0.3,
                rotate_limit=90,
                interpolation=1,
                border_mode=4,
                value=None,
                mask_value=None,
                shift_limit_x=None,
                shift_limit_y=None,
                p=1,
            ),
            A.SafeRotate(
                limit=90,
                interpolation=1,
                border_mode=4,
                value=None,
                mask_value=None,
                always_apply=False,
                p=1,
            ),
        )
    )
    return A.Compose(
        (
            A.HorizontalFlip(p=0.5),
            A.VerticalFlip(p=0.5),
            A.OneOf((A.GaussianBlur(), add_blur, random_shift)),
        ),
        **kwargs,
    )


@cache
def get_trans():
    return define_transformer(bbox_params=A.BboxParams(format="yolo"))


def get_kp(file_path: Path) -> typing.List[tuple]:
    with file_path.open() as fp:
        data = next(fp).split(" ")
        label = [data[0]]
        return [(Stream(data[1:]).map(float).map(lambda x: max(x, 1e-3)).to_list() + label)]


def process(img_path: Path, label_folder: Path, times=3):
    trans = get_trans()
    datasets_folder = img_path.parts[-2]
    label_path = label_folder / datasets_folder / f"{img_path.stem}.txt"
    img = cv2.imdecode(np.fromfile(img_path, dtype=np.uint8), cv2.IMREAD_COLOR)
    kp = get_kp(label_path)

    for _ in range(times):
        ret = trans(image=img, bboxes=kp)
        # save
        file_name = f"{img_path.stem}_aug_{np.random.randint(0, 10000):06}"
        boxes_strs = Stream(ret["bboxes"][0]).map(str).to_list() if ret["bboxes"] else None
        ctx = " ".join((boxes_strs[-1], *boxes_strs[:-1])) if boxes_strs else None

        output_label_file = label_path.parent / f"{file_name}.{label_path.suffix}"
        output_image_file = img_path.parent / f"{file_name}.{img_path.suffix}"
        if ctx is not None and boxes_strs is not None:
            output_label_file.write_text(ctx)
            cv2.imencode(".png", ret["image"])[1].tofile(output_image_file)


def main(imgs: Path, labels: Path, *, debug=False):
    def safe_process(fu: Future):
        if exp := fu.exception():
            print("error: ", exp)
            return

        fu.result()

    pre_process = (
        Stream(imgs.rglob("*"))
        .filter(lambda p: p.suffix in image_suffix)
        .map(lambda i: (i, labels))
    )
    if debug:
        pre_process.map(lambda args: process(*args)).for_each(lambda ret: print(ret))
    else:
        with Exec() as e:
            (
                pre_process.map(lambda args: e.submit(process, *args)).for_each(
                    lambda f: f.add_done_callback(safe_process)
                )
            )


if __name__ == "__main__":
    main(
        Path(r"E:\20240507dataset\datasets\images"),
        Path(r"E:\20240507dataset\datasets\labels"),
        debug=True,
    )
